import pickle

import numpy as np

from alpha_zero.game import Board
from alpha_zero.mcts_alphaZero import MCTSPlayer, MCTS
from alpha_zero.policy_value_net_numpy import PolicyValueNetNumpy
from player.base.player import Player

"""
mcts模型
模型源地址：https://github.com/junxiaosong/AlphaZero_Gomoku.git
"""
class MCTSPlayer(Player):
    model_file = 'best_policy_8_8_5.model'

    temp = 1.0

    def __init__(self, board):
        policy_param = pickle.load(open('alpha_zero/' + self.model_file, 'rb'),
                                   encoding='bytes')
        best_policy = PolicyValueNetNumpy(board.width, board.height, policy_param)
        self.mcts = MCTS(best_policy.policy_value_fn, 5, 400)
        self.board = board
        self._is_selfplay = 0
        self.temp = 1e-3

    def step(self, boardd, pos, code, error_pos) -> (int, int):
        sensible_moves = self.board.availables
        # the pi vector returned by MCTS as in the alphaGo Zero paper
        move_probs = np.zeros(self.board.width * self.board.height)
        if len(sensible_moves) > 0:
            acts, probs = self.mcts.get_move_probs(self.board, self.temp)
            move_probs[list(acts)] = probs
            if self._is_selfplay:
                # add Dirichlet Noise for exploration (needed for
                # self-play training)
                move = np.random.choice(
                    acts,
                    p=0.75 * probs + 0.25 * np.random.dirichlet(0.3 * np.ones(len(probs)))
                )
                # update the root node and reuse the search tree
                self.mcts.update_with_move(move)
            else:
                # with the default temp=1e-3, it is almost equivalent
                # to choosing the move with the highest prob
                move = np.random.choice(acts, p=probs)
                # reset the root node
                self.mcts.update_with_move(-1)
            #                location = board.move_to_location(move)
            #                print("AI move: %d,%d\n" % (location[0], location[1]))

            # print(move)
            location = self.move_to_location(move)
            return int(location[0]),int(location[1])

    def move_to_location(self, move):
        """
        3*3 board's moves like:
        6 7 8
        3 4 5
        0 1 2
        and move 5's location is (1,2)
        """
        h = move // self.board.width
        w = move % self.board.width
        return [h, w]
